##########################################################################################

library(ggplot2)
library(tidyverse)
library(readxl)
library(magrittr)
library(parallel)
library(data.table)
library(patchwork)
library(optparse)
library(rjags)
library(ggmcmc)
library(kableExtra)


set.seed(123)

##########################################################################################

option_list <- list(
    make_option(c("--input_file"), type = "character") ,
    make_option(c("--code_path"), type = "character") ,
    make_option(c("--plot_id_file"), type = "character") ,
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
  #  input_file <- "~/20220915_gastric_multiple/dna_combinePublic/scripts/evolutionTime/ipmn-timing-master/data/Timing_Metrics.xlsx"
  #  out_path <- "~/20220915_gastric_multiple/dna_combinePublic/scripts/evolutionTime"

  input_file <- "~/20220915_gastric_multiple/dna_combinePublic/images/evolutionTime/sample_mutNum.tsv"
  out_path <- "~/20220915_gastric_multiple/dna_combinePublic/images/evolutionTime"
  code_path <- "~/20220915_gastric_multiple/dna_combinePublic/scripts/evolutionTime/ipmn-timing-master"
  plot_id_file <- "~/20220915_gastric_multiple/dna_combinePublic/config/plotID.list"

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

input_file <- opt$input_file
out_path <- opt$out_path
code_path <- opt$code_path
plot_id_file <- opt$plot_id_file

dir.create(out_path , recursive = T)

##########################################################################################

source(paste0(code_path, "/code/functions.R"))
jag_file <- paste0( code_path , "/code/clock.jag"  )

##########################################################################################
dat_plotid <- fread(plot_id_file)

dat <- fread(input_file) %>%
mutate( additional = GC_private ) %>%
mutate( case = ID ) %>%
mutate( high_mutation_rate = additional > 20) %>%
mutate( jags_id=paste0("T[", seq_len(nrow(.)), "]"))

dat$laruen <- ""
dat$laruen <- ifelse( dat$Type == "IM + IGC" , "IGC" , dat$laruen)
dat$laruen <- ifelse( dat$Type == "IM + DGC" , "DGC" , dat$laruen)
dat$laruen <- ifelse( dat$Type == "IM + IGC + DGC(IGC)" , "combine_IGC" , dat$laruen)
dat$laruen <- ifelse( dat$Type == "IM + IGC + DGC(DGC)" , "combine_DGC" , dat$laruen)

dat$case <- paste0( dat$case , "_" , dat$laruen )

y <- dat$additional

##########################################################################################
## mutations/year
## 1-10
result <- Reduce(function(x,y)bind_rows(x,y) , mclapply( seq(1,10,0.01) , function(mu1){

  m1 <- jags.model(jag_file,
                  data=list(y=y, N = length(y), mu=mu1),
                  inits=list(T=rep(5, length(y))),
                  n.adapt = 2000, n.chains = 3)
  jags_samples <- coda.samples(m1,
                               variable.names="T",
                               n.iter=2000)
  jags_df <- ggs(jags_samples)
  
  summaries.mu1 <- jags_df %>%
    "["(grep("T", .$Parameter), ) %>%
    group_by(Parameter) %>%
    summarize(lower=quantile(value, 0.025),
              median=quantile(value, 0.5),
              mean=mean(value),
              upper=quantile(value, 0.975)) %>%
    mutate(mu=mu1)

  return( data.frame(summaries.mu1) )

},mc.cores=30))

result2 <- pivot_longer(result, 
  cols = starts_with(c("lower" , "median" , "mean" , "upper")), 
  names_to = "quantile", 
  values_to = "value")

result2 <- result2[,c("Parameter" , "quantile" , "value" , "mu")]
names( result2 ) <- c("patient" , "quantile" , "value" , "mut.per.yr")

#posteriors <- readRDS(paste0(code_path, "/output/timing_master.R/timing.rds")) %>%
posteriors <- result2 %>%
  mutate(patient=as.character(patient)) %>%
  left_join(select(dat, jags_id, case, high_mutation_rate),
            by=c("patient"="jags_id"))

high.mr <- filter(posteriors, high_mutation_rate, mut.per.yr > 3)
low.mr  <- filter(posteriors, !high_mutation_rate, mut.per.yr < 7)

##########################################################################################
## 长型转宽型
high.mr3 <- high.mr %>%
  pivot_wider(names_from = quantile, values_from = value)
high.mr3 <- high.mr3[,c("case" , "mut.per.yr" , "median" , "lower" , "upper")]

low.mr2 <- low.mr %>%
  pivot_wider(names_from = quantile, values_from = value)
low.mr2 <- low.mr2[,c("case" , "mut.per.yr" , "median" , "lower" , "upper")]

##########################################################################################
post <- bind_rows(low.mr2, high.mr3) %>%
  mutate(case=factor(case, levels=dat$case))

post2 <- group_by(post, case) %>%
  summarize(lower=quantile(median, 0.025),
            mean=mean(median),
            upper=quantile(median, 0.975)) %>%
  mutate(x=11)
one.number <- mean(post2$mean)

##########################################################################################

plotFun <- function(post = post , post2 = post2 , ylab = ylab){
  p <- ggplot(post,
       aes(ymin=lower, ymax=upper,
           y=median, x=mut.per.yr)) +
  geom_ribbon(fill="lightblue", color="lightblue") +
  geom_errorbar(data=post2, aes(ymin=lower, ymax=upper, x=x),
                inherit.aes=FALSE) +
  geom_point(data=post2, aes(x=x, y=mean), inherit.aes=FALSE,
             shape=21, fill="white") +
  geom_line() +
  facet_wrap(~PlotID_Divide) +
  coord_flip() +
  theme(panel.background=element_rect(fill="white", color="black"),
        axis.text=element_text(size=11),
        axis.title=element_text(size=13),
        strip.background=element_rect(fill="white"),
        strip.text=element_text(size=13)) +
  scale_x_continuous(breaks=c(1, 3, 5, 7, 9),
                     labels=c(1, 3, 5, 7, 9),
                     limits=c(1, 12)) +
  xlab("Mutations/year\n") +
  ylim(c(0, 31)) +
  ylab(ylab) 

  return(p)
}

##########################################################################################

post <- subset(post[grep( "combine" , post$case , invert =T ),])
post$ID <- sapply(strsplit(as.character(post$case) , "_") , "[" , 1)
post$Class <- sapply(strsplit(as.character(post$case) , "_") , "[" , 2)
post <- merge( post , dat_plotid[,c("PlotID_Divide" , "ID")] )

post2 <- subset(post2[grep( "combine" , post2$case , invert =T ),])
post2$ID <- sapply(strsplit(as.character(post2$case) , "_") , "[" , 1)
post2$Class <- sapply(strsplit(as.character(post2$case) , "_") , "[" , 2)
post2 <- merge( post2 , dat_plotid[,c("PlotID_Divide" , "ID")] )

class <- "IGC"
ylab <- paste0("Number of years for IM to " , class , " progression")
p1 <- plotFun(post = subset( post , Class == class ) , post2 = subset( post2 , Class == class ) , ylab = ylab)

class <- "DGC"
ylab <- paste0("Number of years for IM to " , class , " progression")
p2 <- plotFun(post = subset( post , Class == class ) , post2 = subset( post2 , Class == class ) , ylab = ylab)

out_name <- paste0( out_path , "/timing_molecular_clock.pdf" )
ggsave( out_name , p1 + p2 , width=15, height=8)

out_name <- paste0( out_path , "/timing_molecular_clock.tsv" )
write.table( post2 , out_name , row.names = F , quote = F , sep = "\t" )

##########################################################################################
## Average of patient-specific means
average <- round(mean(post2$mean), 1)
